from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from typing import Callable
import time
import json
from models.request_log import RequestLog
from models.user import User
from dao.postgresql_connect import get_db
import jwt
import os
import secrets
from fastapi.responses import JSONResponse
import logging
from datetime import datetime
import pytz

# 配置日志
logger = logging.getLogger(__name__)

# JWT Settings
JWT_SECRET_KEY = os.getenv("JWT_SECRET_KEY", secrets.token_urlsafe(32))
JWT_ALGORITHM = "HS256"

# 敏感信息字段
SENSITIVE_FIELDS = {
    'password', 'token', 'secret', 'credit_card', 
    'access_token', 'refresh_token', 'api_key'
}

# 敏感请求头
SENSITIVE_HEADERS = {
    'authorization', 'cookie', 'proxy-authorization',
    'x-api-key', 'api-key'
}

def get_shanghai_time():
    """获取上海时间"""
    shanghai_tz = pytz.timezone('Asia/Shanghai')
    now = datetime.now(shanghai_tz)
    return now

class RequestLoggingMiddleware(BaseHTTPMiddleware):
    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        start_time = time.time()
        
        # 使用新的会话创建方式
        db = next(get_db())
        log_entry = None
        
        try:
            # 创建日志记录
            log_entry = RequestLog(
                ip_address=request.client.host,
                route=request.url.path,
                method=request.method,
                user_agent=request.headers.get("user-agent")
            )
            
            # 记录查询参数
            log_entry.query_params = dict(request.query_params)
            
            # 记录请求头（排除敏感信息）
            headers = dict(request.headers)
            filtered_headers = {k: v for k, v in headers.items() if k.lower() not in SENSITIVE_HEADERS}
            log_entry.headers = filtered_headers
            
            # 尝试获取用户ID和用户名（如果有认证token）
            auth_header = request.headers.get("authorization")
            if auth_header and auth_header.startswith("Bearer "):
                token = auth_header.split(" ")[1]
                try:
                    payload = jwt.decode(token, JWT_SECRET_KEY, algorithms=[JWT_ALGORITHM])
                    user_id = payload.get("user_id")
                    if user_id:
                        log_entry.user_id = user_id
                        # 获取用户信息
                        user = db.query(User).filter(User.id == user_id).first()
                        if user:
                            log_entry.user_name = user.username
                except jwt.InvalidTokenError:
                    pass
            
            # 记录请求体
            if request.method in ["POST", "PUT", "PATCH"]:
                try:
                    body = await request.json()
                    if isinstance(body, dict):
                        filtered_body = {k: '***' if k.lower() in SENSITIVE_FIELDS else v 
                                      for k, v in body.items()}
                        log_entry.request_body = filtered_body
                except Exception as e:
                    logger.error(f"处理请求体失败: {str(e)}")
                    log_entry.request_body = None
            
            # 处理请求并获取响应
            response = await call_next(request)
            log_entry.response_status = response.status_code
            
            # 记录响应体（仅JSON响应）
            if isinstance(response, JSONResponse):
                try:
                    response_body = response.body.decode()
                    if response_body:
                        response_data = json.loads(response_body)
                        if isinstance(response_data, dict):
                            filtered_response = {k: '***' if k.lower() in SENSITIVE_FIELDS else v 
                                              for k, v in response_data.items()}
                            log_entry.response_body = filtered_response
                except Exception as e:
                    logger.error(f"处理响应体失败: {str(e)}")
                    log_entry.response_body = None
            
            # 计算处理时间
            log_entry.processing_time = int((time.time() - start_time) * 1000)
            
            # 保存日志
            db.add(log_entry)
            db.commit()
            
            return response
            
        except Exception as e:
            logger.error(f"请求日志记录失败: {str(e)}", exc_info=True)
            if log_entry:
                log_entry.error_message = str(e)
                log_entry.response_status = 500
                db.add(log_entry)
                db.commit()
            raise
            
        finally:
            db.close()
